import torch
import torch.optim as optim
import torch.nn as nn
import pandas as pd
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import confusion_matrix, roc_curve, auc, classification_report
import seaborn as sns
import matplotlib.pyplot as plt
from MLP import CardioMLP
from config import data_paths, config, result_path

plt.switch_backend('TKAgg')

X_train = pd.read_csv(data_paths['X_train']).values
y_train = pd.read_csv(data_paths['y_train'])['cardio'].values
X_test = pd.read_csv(data_paths['X_test']).values
y_test = pd.read_csv(data_paths['y_test'])['cardio'].values

print("Original y_train shape:", y_train.shape)
print("Original y_test shape:", y_test.shape)

X_train = torch.FloatTensor(X_train).cuda()
y_train = torch.LongTensor(y_train).cuda()
X_test = torch.FloatTensor(X_test).cuda()
y_test = torch.LongTensor(y_test).cuda()

train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

model = CardioMLP().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.0001)


def plot_confusion_matrix(y_true, y_pred):
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=['Negative', 'Positive'],
                yticklabels=['Negative', 'Positive'])
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.savefig(result_path['Confusion Matrix'])


def train_and_evaluate():
    train_losses = []
    test_losses = []
    accuracies = []

    for epoch in range(110):
        model.train()
        total_train_loss = 0
        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_train_loss += loss.item()
        train_losses.append(total_train_loss / len(train_loader))

        model.eval()
        correct = 0
        total_test_loss = 0
        y_true, y_scores = [], []
        with torch.no_grad():
            for data, target in test_loader:
                output = model(data)
                loss = criterion(output, target)
                total_test_loss += loss.item()
                _, predicted = torch.max(output, 1)
                correct += (predicted == target).sum().item()
                probabilities = torch.softmax(output, dim=1)
                y_true.extend(target.cpu().numpy())
                y_scores.extend(probabilities[:, 1].cpu().numpy())

        test_loss = total_test_loss / len(test_loader)
        test_losses.append(test_loss)

        accuracy = 100 * correct / len(test_dataset)
        accuracies.append(accuracy)

        print(f'Epoch {epoch + 1}: Train Loss: {train_losses[-1]:.4f}, '
              f'Test Loss: {test_losses[-1]:.4f}, '
              f'Accuracy: {accuracy:.2f}%, ')

    torch.save(model.state_dict(), 'cardio_model.pth')

    y_pred = [1 if score > 0.5 else 0 for score in y_scores]
    print(classification_report(y_true, y_pred))
    plot_confusion_matrix(y_true, y_pred)

    fpr, tpr, _ = roc_curve(y_true, y_scores)
    roc_auc = auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver operating characteristic')
    plt.legend(loc="lower right")
    plt.savefig(result_path['ROC'])

    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label='Train Loss')
    plt.plot(test_losses, label='Test Loss')
    plt.title('Training and Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig(result_path['Loss Over Epochs'])

    plt.figure(figsize=(10, 5))
    plt.plot(accuracies, label='Accuracy')
    plt.title('Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.legend()
    plt.savefig(result_path['Accuracy Over Epochs'])


if __name__ == "__main__":
    train_and_evaluate()
